
import com.forever.toolAndAgent.Assistant;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.document.parser.apache.tika.ApacheTikaDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentByParagraphSplitter;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.Tokenizer;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.model.openai.OpenAiTokenizer;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.Result;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.milvus.MilvusEmbeddingStore;
import org.junit.jupiter.api.Test;


import java.util.List;

public class DocumentLoad {
    @Test
    public void LoadAndParser(){
        // Load all documents from a directory
        //可以解析txt，word会乱码
        Document document= FileSystemDocumentLoader.loadDocument("C:/Users/15143/Desktop/新建 文本文档 (2).txt", new TextDocumentParser());
        System.out.println(document);
    }


    @Test
    public void LoadAndParser2(){
        //可以解析word
        Document document= FileSystemDocumentLoader.loadDocument("C:/Users/15143/Desktop/新建 文本文档 (2).txt", new TextDocumentParser());
        Document document1= FileSystemDocumentLoader.loadDocument("C:/Users/15143/Desktop/新建文件夹/新建 DOC 文档.doc", new ApacheTikaDocumentParser());
        System.out.println(document);

    }

    @Test
    public void LoadAndParser3(){
        EmbeddingModel model =
                new OpenAiEmbeddingModel
                .OpenAiEmbeddingModelBuilder()
                        .modelName("quentinz/bge-base-zh-v1.5")
                        .baseUrl("http://localhost:11434/v1")
                        .apiKey("ollama")
                        .dimensions(1536)
                        .build();


        // Create a tokenizer instance
        Tokenizer tokenizer = new OpenAiTokenizer();

        // Create a DocumentSplitter with a max segment size of 1024 tokens
        DocumentByParagraphSplitter splitter = new DocumentByParagraphSplitter(1024, 0, tokenizer);
        Document document = FileSystemDocumentLoader.loadDocument("C:/Users/15143/Desktop/年轻人喜欢养石头宠物.docx", new ApacheTikaDocumentParser());
        List<TextSegment> segments = splitter.split(document);


        EmbeddingStore embeddingStore = MilvusEmbeddingStore.builder().uri("http://localhost:19530")
                .collectionName("test_collection")
                .dimension(768)
                .build();
        for (TextSegment segment : segments) {
            Response<Embedding> embed = model.embed(segment);
            System.out.println(segment);
            System.out.println(embed);
            System.out.println(embed.content().dimension());
            System.out.println();

            embeddingStore.add(embed.content(), segment);
        }



        ChatLanguageModel chatmodel =
                new OpenAiChatModel
                        .OpenAiChatModelBuilder()
                        .baseUrl("http://localhost:11434/v1")
                        .apiKey("ollama")
                        .modelName("qwen2:7b")
                        .logRequests(true)
                        .logResponses(true)
                        .build();

        ContentRetriever contentRetriever = new EmbeddingStoreContentRetriever(embeddingStore, model);

        Assistant assistant = AiServices.builder(Assistant.class)
                .chatLanguageModel(chatmodel)
                .contentRetriever(contentRetriever)
                .build();

        Result answer = assistant.chat("年轻人为什么喜欢石头宠物");
        System.out.println(answer);

    }
}
